import sqlite3
from contextlib import contextmanager
from config import Config
from utils.logger import get_logger

logger = get_logger(__name__)


class Database:
    """数据库管理类"""

    def __init__(self, db_file=None):
        self.db_file = db_file or Config.DATABASE_FILE

    @contextmanager
    def get_connection(self):
        """数据库连接上下文管理器"""
        conn = sqlite3.connect(self.db_file)
        conn.row_factory = sqlite3.Row
        try:
            yield conn
        finally:
            conn.close()

    def init_database(self):
        """初始化数据库表"""
        with self.get_connection() as conn:
            cursor = conn.cursor()

            # 用户上传记录表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS uploads (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    file_id TEXT UNIQUE NOT NULL,
                    original_name TEXT NOT NULL,
                    file_type TEXT NOT NULL,
                    file_path TEXT NOT NULL,
                    preview_path TEXT,
                    upload_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    file_size INTEGER,
                    is_deleted INTEGER DEFAULT 0
                )
            ''')

            # 处理任务记录表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS tasks (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    task_id TEXT UNIQUE NOT NULL,
                    human_file_id TEXT,
                    background_file_id TEXT,
                    parameters TEXT,
                    status TEXT DEFAULT 'pending',
                    progress INTEGER DEFAULT 0,
                    trained INTEGER DEFAULT 0,
                    result_path TEXT,
                    preview_path TEXT,
                    error_message TEXT,
                    created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    completed_time TIMESTAMP,
                    processing_time_seconds INTEGER,
                    output_name TEXT,
                    is_deleted INTEGER DEFAULT 0
                )
            ''')

            # 预览记录表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS previews (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    preview_id TEXT UNIQUE NOT NULL,
                    human_file_id TEXT,
                    background_file_id TEXT,
                    preview_path TEXT,
                    parameters TEXT,
                    created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')

            # 预设表
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS presets (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    name TEXT NOT NULL,
                    description TEXT,
                    canvas_width INTEGER,
                    canvas_height INTEGER,
                    output_format TEXT,
                    video_scale REAL,
                    video_x INTEGER,
                    video_y INTEGER,
                    bg_scale REAL,
                    bg_x INTEGER,
                    bg_y INTEGER,
                    keep_human_ratio INTEGER,
                    keep_bg_ratio INTEGER,
                    created_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            ''')

            conn.commit()
            logger.info("数据库初始化完成")


class UploadDAO:
    """上传文件数据访问对象"""

    def __init__(self, db: Database):
        self.db = db

    def save_upload(self, file_id, original_name, file_type, file_path,
                    preview_path=None, file_size=0):
        """保存上传记录"""
        try:
            with self.db.get_connection() as conn:
                conn.execute('''
                    INSERT OR REPLACE INTO uploads 
                    (file_id, original_name, file_type, file_path, preview_path, file_size)
                    VALUES (?, ?, ?, ?, ?, ?)
                ''', (file_id, original_name, file_type, file_path, preview_path, file_size))
                conn.commit()
                return True
        except Exception as e:
            logger.error(f"保存上传记录失败: {e}")
            return False

    def get_upload_by_id(self, file_id):
        """根据ID获取上传记录"""
        with self.db.get_connection() as conn:
            return conn.execute('''
                SELECT * FROM uploads WHERE file_id = ? AND is_deleted = 0
            ''', (file_id,)).fetchone()

    def get_uploads(self, file_type='all', page=1, page_size=12, search='',
                    sort_by='time_desc', filter_type='all'):
        """获取上传列表"""
        with self.db.get_connection() as conn:
            query = 'SELECT * FROM uploads WHERE is_deleted = 0'
            params = []

            if file_type != 'all':
                query += ' AND file_type = ?'
                params.append(file_type)

            if search:
                query += ' AND original_name LIKE ?'
                params.append(f'%{search}%')

            if filter_type != 'all' and file_type == 'background':
                if filter_type == 'image':
                    query += " AND (original_name LIKE '%.jpg' OR original_name LIKE '%.jpeg' OR original_name LIKE '%.png')"
                elif filter_type == 'gif':
                    query += " AND original_name LIKE '%.gif'"
                elif filter_type == 'video':
                    query += " AND (original_name LIKE '%.mp4' OR original_name LIKE '%.mov')"

            # 获取总数
            count_query = query.replace('SELECT *', 'SELECT COUNT(*) as count')
            total = conn.execute(count_query, params).fetchone()['count']

            # 排序
            order_by = {
                'time_desc': ' ORDER BY upload_time DESC',
                'time_asc': ' ORDER BY upload_time ASC',
                'name': ' ORDER BY original_name ASC',
                'size': ' ORDER BY file_size DESC'
            }.get(sort_by, ' ORDER BY upload_time DESC')

            query += order_by

            # 分页
            offset = (page - 1) * page_size
            query += ' LIMIT ? OFFSET ?'
            params.extend([page_size, offset])

            files = conn.execute(query, params).fetchall()

            return {
                'files': files,
                'total': total,
                'page': page,
                'page_size': page_size
            }

    def delete_upload(self, file_id):
        """软删除上传记录"""
        try:
            with self.db.get_connection() as conn:
                conn.execute('''
                    UPDATE uploads SET is_deleted = 1 WHERE file_id = ?
                ''', (file_id,))
                conn.commit()
                return True
        except Exception as e:
            logger.error(f"删除上传记录失败: {e}")
            return False

    def batch_delete_uploads(self, file_ids):
        """批量删除上传记录"""
        try:
            with self.db.get_connection() as conn:
                placeholders = ','.join(['?' for _ in file_ids])
                conn.execute(f'''
                    UPDATE uploads SET is_deleted = 1 
                    WHERE file_id IN ({placeholders})
                ''', file_ids)
                conn.commit()
                return True
        except Exception as e:
            logger.error(f"批量删除上传记录失败: {e}")
            return False


class TaskDAO:
    """任务数据访问对象"""

    def __init__(self, db: Database):
        self.db = db

    def create_task(self, task_id, human_file_id, background_file_id,
                    parameters, output_name, status='pending'):
        """创建任务"""
        try:
            with self.db.get_connection() as conn:
                conn.execute('''
                    INSERT INTO tasks 
                    (task_id, human_file_id, background_file_id, parameters, status, output_name)
                    VALUES (?, ?, ?, ?, ?, ?)
                ''', (task_id, human_file_id, background_file_id, parameters, status, output_name))
                conn.commit()
                return True
        except Exception as e:
            logger.error(f"创建任务失败: {e}")
            return False

    def update_task_status(self, task_id, status, progress=None, error_message=None):
        """更新任务状态"""
        try:
            with self.db.get_connection() as conn:
                if progress is not None:
                    conn.execute('''
                        UPDATE tasks SET status = ?, progress = ?
                        WHERE task_id = ?
                    ''', (status, progress, task_id))
                else:
                    conn.execute('''
                        UPDATE tasks SET status = ?, error_message = ?
                        WHERE task_id = ?
                    ''', (status, error_message, task_id))
                conn.commit()
                return True
        except Exception as e:
            logger.error(f"更新任务状态失败: {e}")
            return False

    def complete_task(self, task_id, result_path, preview_path, processing_time):
        """完成任务"""
        try:
            with self.db.get_connection() as conn:
                conn.execute('''
                    UPDATE tasks SET 
                        status = 'completed',
                        result_path = ?,
                        preview_path = ?,
                        completed_time = CURRENT_TIMESTAMP,
                        processing_time_seconds = ?
                    WHERE task_id = ?
                ''', (result_path, preview_path, processing_time, task_id))
                conn.commit()
                return True
        except Exception as e:
            logger.error(f"完成任务失败: {e}")
            return False

    def get_task_by_id(self, task_id):
        """根据ID获取任务"""
        with self.db.get_connection() as conn:
            return conn.execute('''
                SELECT * FROM tasks WHERE task_id = ? AND is_deleted = 0
            ''', (task_id,)).fetchone()

    def get_task_history(self, page=1, page_size=10, search='', status='', start_date='', end_date=''):
        """获取任务历史记录"""
        with self.db.get_connection() as conn:
            offset = (page - 1) * page_size

            # 构建查询条件
            conditions = []
            params = []

            if search:
                conditions.append('(t.output_name LIKE ? OR u1.original_name LIKE ? OR u2.original_name LIKE ?)')
                search_param = f'%{search}%'
                params.extend([search_param, search_param, search_param])

            if status:
                # 处理特殊状态：completed但未训练
                if status == 'completed':
                    conditions.append("(t.status = 'completed' AND (t.trained IS NULL OR t.trained = 0))")
                else:
                    conditions.append("(t.status = 'completed' AND t.trained = ?)")
                    params.append(1)

            if start_date:
                conditions.append('DATE(t.created_time) >= ?')
                params.append(start_date)

            if end_date:
                conditions.append('DATE(t.completed_time) <= ?')
                params.append(end_date)

            where_clause = 'WHERE ' + ' AND '.join(conditions) if conditions else ''

            # 查询总数
            count_query = f'''
                SELECT COUNT(*) as total
                FROM tasks t
                LEFT JOIN uploads u1 ON t.human_file_id = u1.file_id
                LEFT JOIN uploads u2 ON t.background_file_id = u2.file_id
                {where_clause}
            '''

            total = conn.execute(count_query, params).fetchone()['total']

            # 查询数据
            data_query = f'''
                SELECT 
                    t.*,
                    u1.original_name as human_name,
                    u2.original_name as background_name,
                    COALESCE(t.trained, 0) as trained
                FROM tasks t
                LEFT JOIN uploads u1 ON t.human_file_id = u1.file_id
                LEFT JOIN uploads u2 ON t.background_file_id = u2.file_id
                {where_clause}
                ORDER BY t.created_time DESC
                LIMIT ? OFFSET ?
            '''

            params.extend([page_size, offset])
            tasks = conn.execute(data_query, params).fetchall()

            return {
                'tasks': [dict(row) for row in tasks],
                'total': total,
                'page': page,
                'page_size': page_size
            }

    def delete_task(self, task_id):
        """删除任务"""
        try:
            with self.db.get_connection() as conn:
                conn.execute('''
                    DELETE FROM tasks WHERE task_id = ?
                ''', (task_id,))
                conn.commit()
                return True
        except Exception as e:
            logger.error(f"删除任务失败: {e}")
            return False

    def update_task_trained_status(self, task_id, trained=True):
        """更新任务的训练状态"""
        with self.db.get_connection() as conn:
            conn.execute('''
                UPDATE tasks 
                SET trained = ?, completed_time = CURRENT_TIMESTAMP
                WHERE task_id = ?
            ''', (1 if trained else 0, task_id))
            conn.commit()



class PreviewDAO:
    """预览数据访问对象"""

    def __init__(self, db: Database):
        self.db = db

    def save_preview(self, preview_id, human_file_id, background_file_id,
                     preview_path, parameters):
        """保存预览记录"""
        try:
            with self.db.get_connection() as conn:
                conn.execute('''
                    INSERT INTO previews 
                    (preview_id, human_file_id, background_file_id, preview_path, parameters)
                    VALUES (?, ?, ?, ?, ?)
                ''', (preview_id, human_file_id, background_file_id, preview_path, parameters))
                conn.commit()
                return True
        except Exception as e:
            logger.error(f"保存预览记录失败: {e}")
            return False

    def delete_preview(self, preview_id):
        """删除预览记录"""
        try:
            with self.db.get_connection() as conn:
                conn.execute('DELETE FROM previews WHERE preview_id = ?', (preview_id,))
                conn.commit()
                return True
        except Exception as e:
            logger.error(f"删除预览记录失败: {e}")
            return False


class PresetDAO:
    """预设数据访问对象"""

    def __init__(self, db: Database):
        self.db = db

    def get_presets(self):
        """获取所有预设"""
        with self.db.get_connection() as conn:
            return conn.execute('SELECT * FROM presets ORDER BY created_time DESC').fetchall()

    def save_preset(self, preset_data):
        """保存预设"""
        try:
            with self.db.get_connection() as conn:
                if preset_data.get('id'):
                    # 更新现有预设
                    conn.execute('''
                        UPDATE presets SET 
                            name = ?, description = ?, canvas_width = ?, canvas_height = ?,
                            output_format = ?, video_scale = ?, video_x = ?, video_y = ?,
                            bg_scale = ?, bg_x = ?, bg_y = ?, keep_human_ratio = ?, 
                            keep_bg_ratio = ?, updated_time = CURRENT_TIMESTAMP
                        WHERE id = ?
                    ''', (
                        preset_data['name'], preset_data.get('description', ''),
                        preset_data['canvas_width'], preset_data['canvas_height'],
                        preset_data['output_format'], preset_data['video_scale'],
                        preset_data.get('video_x', 0), preset_data.get('video_y', 0),
                        preset_data['bg_scale'], preset_data.get('bg_x', 0),
                        preset_data.get('bg_y', 0), int(preset_data['keep_human_ratio']),
                        int(preset_data['keep_bg_ratio']), preset_data['id']
                    ))
                else:
                    # 创建新预设
                    conn.execute('''
                        INSERT INTO presets (
                            name, description, canvas_width, canvas_height,
                            output_format, video_scale, video_x, video_y,
                            bg_scale, bg_x, bg_y, keep_human_ratio, keep_bg_ratio, created_time
                        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, current_timestamp)
                    ''', (
                        preset_data['name'], preset_data.get('description', ''),
                        preset_data['canvas_width'], preset_data['canvas_height'],
                        preset_data['output_format'], preset_data['video_scale'],
                        preset_data.get('video_x', 0), preset_data.get('video_y', 0),
                        preset_data['bg_scale'], preset_data.get('bg_x', 0),
                        preset_data.get('bg_y', 0), int(preset_data['keep_human_ratio']),
                        int(preset_data['keep_bg_ratio'])
                    ))
                conn.commit()
                return True
        except Exception as e:
            logger.error(f"保存预设失败: {e}")
            return False

    def delete_preset(self, preset_id):
        """删除预设"""
        try:
            with self.db.get_connection() as conn:
                conn.execute('DELETE FROM presets WHERE id = ?', (preset_id,))
                conn.commit()
                return True
        except Exception as e:
            logger.error(f"删除预设失败: {e}")
            return False